package top.yudoge.dogeeditor.llm.memory

import org.springframework.stereotype.Component
import top.yudoge.dogeeditor.commons.cache.Cache
import top.yudoge.dogeeditor.commons.cache.InMemoryCache
import java.util.HashMap

/**
 * An `SessionMemory` implementation based on `Cache`
 */
@Component
class CacheSessionMemory : SessionMemory {

    private val cache: Cache<String> = InMemoryCache() // we use in-memory cache for now... QAQ!!!!

    override fun newSession(sessionId: String, expireTimeInMS: Long): Boolean {
        if (cache.contains(memoryKey(sessionId))) {
            return false
        }
        cache.putWithExpireTime(memoryKey(sessionId), HashMap<String, String>(), expireTimeInMS)
        return true
    }

    override fun renewSession(sessionId: String, expireTimeInMS: Long): Boolean {
        val key = memoryKey(sessionId)
        if (!cache.contains(key)) {
            return false
        }
        var memory = getMemory(sessionId)!!
        cache.putWithExpireTime(key, memory, expireTimeInMS)
        return true
    }

    override fun get(sessionId: String, key: String): String? {
        return getMemory(sessionId)?.get(key)
    }

    override fun put(sessionId: String, key: String, value: String): Boolean {
        val memory = getMemory(sessionId) ?: return false
        memory[key] = value
        cache.put(memoryKey(sessionId), memory)
        return true
    }

    override fun contains(sessionId: String, key: String): Boolean {
        return getMemory(sessionId)?.containsKey(key) ?: false
    }

    override fun containsSession(sessionId: String): Boolean {
        return cache.contains(memoryKey(sessionId))
    }

    override fun remove(sessionId: String, key: String): Boolean {
        val memory = getMemory(sessionId) ?: return false
        memory.remove(key)
        return true
    }

    override fun revoke(sessionId: String): Boolean {
        val key = memoryKey(sessionId)
        if (!cache.contains(key)) {
            return false
        }
        cache.remove(key)
        return true
    }

    private fun getMemory(sessionId: String): HashMap<String, String>? {
        val memory = cache.get(memoryKey(sessionId), HashMap::class.java)?.let {
            println(it)
            it as HashMap<String, String>
        }
        return memory
    }

    private inline fun memoryKey(sessionId: String): String = "session:memory:$sessionId"

}